"""
现代化Web UI主应用
"""
from flask import Flask, render_template, jsonify, request
from flask_socketio import SocketIO, emit
import sys
import os
from pathlib import Path
import json
import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.utils
import numpy as np
from datetime import datetime, timedelta
import random

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

try:
    from src.data_storage import DataStorage
    from src.data_source_manager import DataSourceManager
    from src.realtime_data import RealTimeDataFetcher
except ImportError:
    # 如果src模块导入失败，设置为None
    DataStorage = None
    DataSourceManager = None
    RealTimeDataFetcher = None

app = Flask(__name__)
app.config['SECRET_KEY'] = 'quantitative_stock_analysis_2024'
socketio = SocketIO(app, cors_allowed_origins="*")

# 全局变量
data_storage = None
data_source_manager = None
realtime_fetcher = None

def init_components():
    """初始化组件"""
    global data_storage, data_source_manager, realtime_fetcher
    try:
        if DataStorage:
            data_storage = DataStorage()
        if DataSourceManager:
            data_source_manager = DataSourceManager()
        if RealTimeDataFetcher:
            realtime_fetcher = RealTimeDataFetcher()
        print("✅ Web UI组件初始化成功")
    except Exception as e:
        print(f"❌ Web UI组件初始化失败: {e}")

@app.route('/')
def index():
    """主页"""
    return render_template('index.html')

@app.route('/chart')
def chart():
    """专业K线图页面"""
    return render_template('chart.html')

@app.route('/api/test_stock/<symbol>')
def get_test_stock_data(symbol):
    """获取测试股票数据（直接调用模拟数据生成器）"""
    try:
        period = request.args.get('period', '1d')
        days = int(request.args.get('days', 60))
        
        # 直接生成模拟数据
        data = generate_mock_stock_data(symbol, period, days)
        if data is None:
            return jsonify({'error': f'无法生成股票 {symbol} 的数据'}), 404
            
        # 转换为JSON格式
        data_json = data.reset_index().to_dict('records')
        
        # 转换日期格式
        for record in data_json:
            if 'Date' in record:
                record['Date'] = record['Date'].strftime('%Y-%m-%d')
            elif 'date' in record:
                record['date'] = record['date'].strftime('%Y-%m-%d')
                
        return jsonify({
            'symbol': symbol,
            'name': get_stock_name(symbol),
            'period': period,
            'data': data_json
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/api/stocks')
def get_stocks():
    """获取股票列表"""
    try:
        if not data_storage:
            return jsonify({'error': '数据存储未初始化'}), 500
            
        # 获取可用股票列表
        stocks = []
        data_dir = Path('data/stocks/daily')
        if data_dir.exists():
            for file in data_dir.glob('*_1d.csv'):
                symbol = file.stem.replace('_1d', '')
                if len(symbol) == 6 and symbol.isdigit():
                    # 尝试获取股票名称
                    name = get_stock_name(symbol)
                    stocks.append({
                        'symbol': symbol,
                        'name': name,
                        'display': f"{symbol} - {name}"
                    })
        
        # 按股票代码排序
        stocks.sort(key=lambda x: x['symbol'])
        return jsonify(stocks)
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/stock/<symbol>')
def get_stock_data(symbol):
    """获取股票数据"""
    try:
        period = request.args.get('period', '1d')
        days = int(request.args.get('days', 60))
        
        # 加载股票数据
        data = load_stock_data(symbol, period, days)
        if data is None:
            return jsonify({'error': f'无法加载股票 {symbol} 的数据'}), 404
            
        # 转换为JSON格式
        data_json = data.reset_index().to_dict('records')
        
        # 转换日期格式
        for record in data_json:
            if 'Date' in record:
                record['Date'] = record['Date'].strftime('%Y-%m-%d')
            elif 'date' in record:
                record['date'] = record['date'].strftime('%Y-%m-%d')
                
        return jsonify({
            'symbol': symbol,
            'name': get_stock_name(symbol),
            'period': period,
            'data': data_json
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/kline/<symbol>')
def get_kline_chart(symbol):
    """生成K线图"""
    try:
        period = request.args.get('period', '1d')
        days = int(request.args.get('days', 60))
        
        # 加载股票数据
        data = load_stock_data(symbol, period, days)
        if data is None:
            return jsonify({'error': f'无法加载股票 {symbol} 的数据'}), 404
            
        # 创建K线图
        fig = create_kline_chart(data, symbol, period)
        
        # 转换为JSON
        graphJSON = json.dumps(fig, cls=plotly.utils.PlotlyJSONEncoder)
        
        return jsonify({
            'symbol': symbol,
            'name': get_stock_name(symbol),
            'period': period,
            'chart': graphJSON
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

def generate_mock_stock_data(symbol: str, period: str = "1d", days: int = 60):
    """生成模拟股票数据"""
    try:
        # 初始价格设置
        base_prices = {
            '000001': 13.0,  # 平安银行
            '000002': 18.5,  # 万科A
            '000858': 185.0,  # 五粮液
        }
        base_price = base_prices.get(symbol, 50.0)
        
        # 生成日期序列
        end_date = datetime.now().date()
        start_date = end_date - timedelta(days=days-1)
        date_range = pd.date_range(start=start_date, end=end_date, freq='D')
        
        # 初始化数据
        data = []
        current_price = base_price
        
        # 设置随机种子以获得一致的结果
        np.random.seed(hash(symbol) % 1000)
        random.seed(hash(symbol) % 1000)
        
        for i, date in enumerate(date_range):
            # 模拟更真实的股价波动
            # 使用正弦波和随机波动结合
            trend = np.sin(i * 0.1) * 0.02  # 长期趋势
            volatility = np.random.normal(0, 0.015)  # 随机波动
            daily_change = trend + volatility
            
            # 计算当天的OHLC
            open_price = current_price * (1 + np.random.normal(0, 0.005))
            
            # 日内波动范围
            intraday_range = abs(np.random.normal(0, 0.02))
            high_offset = np.random.uniform(0.3, 1.0) * intraday_range
            low_offset = np.random.uniform(0.3, 1.0) * intraday_range
            
            high_price = open_price * (1 + high_offset)
            low_price = open_price * (1 - low_offset)
            
            # 收盘价基于当日趋势
            close_direction = np.random.choice([-1, 1], p=[0.45, 0.55])  # 略偏上涨
            close_range = np.random.uniform(0.2, 0.8)
            if close_direction > 0:
                close_price = low_price + (high_price - low_price) * close_range
            else:
                close_price = high_price - (high_price - low_price) * close_range
            
            # 确保OHLC逻辑正确
            high_price = max(high_price, open_price, close_price)
            low_price = min(low_price, open_price, close_price)
            
            # 生成成交量（与价格波动成反比）
            price_volatility = abs(high_price - low_price) / open_price
            base_volume = 100000000  # 1亿
            volume = int(base_volume * (1 + price_volatility * 2) * np.random.uniform(0.5, 1.5))
            
            data.append({
                'Date': date,
                'open': round(open_price, 2),
                'high': round(high_price, 2),
                'low': round(low_price, 2),
                'close': round(close_price, 2),
                'volume': volume
            })
            
            # 更新当前价格供下一天使用
            current_price = close_price
        
        # 创建DataFrame
        df = pd.DataFrame(data)
        df.set_index('Date', inplace=True)
        
        # 如果是周线或月线，进行聚合
        if period == '1w':
            df = df.resample('W').agg({
                'open': 'first',
                'high': 'max',
                'low': 'min',
                'close': 'last',
                'volume': 'sum'
            }).dropna()
        elif period == '1M':
            df = df.resample('M').agg({
                'open': 'first',
                'high': 'max',
                'low': 'min',
                'close': 'last',
                'volume': 'sum'
            }).dropna()
        
        return df
        
    except Exception as e:
        print(f"生成模拟数据失败: {e}")
        return None

def load_stock_data(symbol: str, period: str = "1d", days: int = 60):
    """加载股票数据"""
    try:
        if not data_storage:
            # 如果没有数据存储，生成模拟数据
            return generate_mock_stock_data(symbol, period, days)
            
        # 构建文件路径
        if period in ['1d', '1w', '1M']:  # 日线、周线、月线都使用日线数据
            file_path = Path(f'data/stocks/daily/{symbol}_1d.csv')
        elif period == "1h":
            file_path = Path(f'data/stocks/hourly/{symbol}_1h.csv')
        else:
            file_path = Path(f'data/stocks/minute/{symbol}_{period}.csv')
            
        if not file_path.exists():
            # 如果文件不存在，生成模拟数据
            return generate_mock_stock_data(symbol, period, days)
            
        # 读取CSV文件
        df = pd.read_csv(file_path)
        
        # 处理日期列
        if 'Date' in df.columns:
            df['Date'] = pd.to_datetime(df['Date'])
            df.set_index('Date', inplace=True)
        elif 'date' in df.columns:
            df['date'] = pd.to_datetime(df['date'])
            df.set_index('date', inplace=True)
        else:
            df.index = pd.to_datetime(df.index)
            
        # 标准化列名
        df.columns = df.columns.str.lower()

        # 确保必要的列存在
        required_columns = ['open', 'high', 'low', 'close', 'volume']
        if not all(col in df.columns for col in required_columns):
            return None

        # 确保数据类型正确
        for col in ['open', 'high', 'low', 'close', 'volume']:
            df[col] = pd.to_numeric(df[col], errors='coerce')

        # 删除包含NaN的行
        df = df.dropna(subset=['open', 'high', 'low', 'close'])
            
        # 如果是周线或月线，需要进行数据聚合
        if period == '1w':  # 周线
            # 重采样为周数据
            df = df.resample('W').agg({
                'open': 'first',
                'high': 'max',
                'low': 'min',
                'close': 'last',
                'volume': 'sum'
            })
            df = df.dropna()
            days = min(days // 5, len(df))  # 调整天数
        elif period == '1M':  # 月线
            # 重采样为月数据
            df = df.resample('M').agg({
                'open': 'first',
                'high': 'max',
                'low': 'min',
                'close': 'last',
                'volume': 'sum'
            })
            df = df.dropna()
            days = min(days // 20, len(df))  # 调整天数
        
        # 获取最近的数据
        if len(df) > days:
            df = df.tail(days)
            
        return df
        
    except Exception as e:
        print(f"加载股票数据失败: {e}")
        return None

def get_stock_name(symbol: str) -> str:
    """获取股票名称"""
    # 使用简单的股票名称映射
    stock_names = {
        '000001': '平安银行',
        '000002': '万科A',
        '000858': '五粮液',
        '002415': '海康威视',
        '002594': '比亚迪',
        '600000': '浦发银行',
        '600036': '招商银行',
        '600519': '贵州茅台',
        '601318': '中国平安',
        '601398': '工商银行'
    }
    return stock_names.get(symbol, f'股票{symbol}')

@app.route('/api/realtime/<symbol>')
def get_realtime_data(symbol):
    """获取实时行情数据"""
    try:
        if realtime_fetcher:
            data = realtime_fetcher.get_realtime_quote(symbol)
            return jsonify(data)
        else:
            return jsonify({'error': '实时数据服务未启用'}), 503
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/realtime/batch', methods=['POST'])
def get_batch_realtime():
    """批量获取实时行情"""
    try:
        if not realtime_fetcher:
            return jsonify({'error': '实时数据服务未启用'}), 503
            
        symbols = request.json.get('symbols', [])
        if not symbols:
            return jsonify({'error': '未提供股票代码'}), 400
            
        data = realtime_fetcher.get_batch_quotes(symbols)
        return jsonify(data)
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/market/overview')
def get_market_overview():
    """获取市场概览"""
    try:
        if realtime_fetcher:
            data = realtime_fetcher.get_market_overview()
            return jsonify(data)
        else:
            return jsonify({'error': '实时数据服务未启用'}), 503
    except Exception as e:
        return jsonify({'error': str(e)}), 500

def create_kline_chart(data: pd.DataFrame, symbol: str, period: str):
    """创建K线图"""
    # 计算技术指标
    data = calculate_indicators(data.copy())
    
    # 创建子图
    fig = make_subplots(
        rows=4, cols=1,
        shared_xaxes=True,
        vertical_spacing=0.03,
        subplot_titles=('K线图', '成交量', 'MACD', 'RSI'),
        row_heights=[0.5, 0.2, 0.15, 0.15]
    )
    
    # 添加蜡烛图
    fig.add_trace(
        go.Candlestick(
            x=data.index,
            open=data['open'],
            high=data['high'],
            low=data['low'],
            close=data['close'],
            name='蜡烛图',
            increasing_line_color='#ef4444',  # 红色上涨
            decreasing_line_color='#22c55e',  # 绿色下跌
            increasing_fillcolor='#ef4444',
            decreasing_fillcolor='#22c55e',
            line=dict(width=1),
            showlegend=True
        ),
        row=1, col=1
    )
    
    # 添加移动平均线
    fig.add_trace(go.Scatter(x=data.index, y=data['MA5'], mode='lines', name='MA5', line=dict(color='#3b82f6', width=2)), row=1, col=1)
    fig.add_trace(go.Scatter(x=data.index, y=data['MA10'], mode='lines', name='MA10', line=dict(color='#8b5cf6', width=2)), row=1, col=1)
    fig.add_trace(go.Scatter(x=data.index, y=data['MA20'], mode='lines', name='MA20', line=dict(color='#f59e0b', width=2)), row=1, col=1)
    
    # 添加成交量
    colors = ['#ef4444' if close >= open_price else '#22c55e' for close, open_price in zip(data['close'], data['open'])]
    fig.add_trace(go.Bar(x=data.index, y=data['volume'], name='成交量', marker_color=colors, opacity=0.7), row=2, col=1)

    # 添加MACD
    fig.add_trace(go.Scatter(x=data.index, y=data['MACD'], mode='lines', name='MACD', line=dict(color='#3b82f6', width=2)), row=3, col=1)
    fig.add_trace(go.Scatter(x=data.index, y=data['MACD_signal'], mode='lines', name='Signal', line=dict(color='#f59e0b', width=2)), row=3, col=1)

    colors = ['#ef4444' if h >= 0 else '#22c55e' for h in data['MACD_histogram']]
    fig.add_trace(go.Bar(x=data.index, y=data['MACD_histogram'], name='MACD柱', marker_color=colors, opacity=0.7), row=3, col=1)

    # 添加RSI
    fig.add_trace(go.Scatter(x=data.index, y=data['RSI'], mode='lines', name='RSI', line=dict(color='#8b5cf6', width=2)), row=4, col=1)
    fig.add_hline(y=70, line_dash="dash", line_color="#ef4444", row=4, col=1)
    fig.add_hline(y=30, line_dash="dash", line_color="#22c55e", row=4, col=1)
    fig.add_hline(y=50, line_dash="solid", line_color="#6b7280", opacity=0.5, row=4, col=1)
    
    # 设置布局
    period_name = {"1d": "日线", "1h": "小时线", "30m": "30分钟", "15m": "15分钟", "5m": "5分钟"}.get(period, "日线")
    
    fig.update_layout(
        title=f"{symbol} - {get_stock_name(symbol)} ({period_name})",
        template="plotly_white",
        showlegend=True,
        height=800,
        hovermode='x unified',
        paper_bgcolor='white',
        plot_bgcolor='white',
        font=dict(color='#374151'),
        title_font=dict(color='#1f2937', size=18)
    )
    
    fig.update_xaxes(rangeslider_visible=False)
    fig.update_yaxes(range=[0, 100], row=4, col=1)
    
    return fig

def calculate_indicators(data: pd.DataFrame) -> pd.DataFrame:
    """计算技术指标"""
    # 移动平均线
    data['MA5'] = data['close'].rolling(window=5).mean()
    data['MA10'] = data['close'].rolling(window=10).mean()
    data['MA20'] = data['close'].rolling(window=20).mean()
    
    # MACD
    exp1 = data['close'].ewm(span=12).mean()
    exp2 = data['close'].ewm(span=26).mean()
    data['MACD'] = exp1 - exp2
    data['MACD_signal'] = data['MACD'].ewm(span=9).mean()
    data['MACD_histogram'] = data['MACD'] - data['MACD_signal']
    
    # RSI
    delta = data['close'].diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
    rs = gain / loss
    data['RSI'] = 100 - (100 / (1 + rs))
    
    return data

if __name__ == '__main__':
    init_components()
    socketio.run(app, debug=True, host='127.0.0.1', port=5000)
